from collections.abc import Sequence
from functools import cached_property
from typing import Callable

import numpy as np
import pyzstd
from zkl_pyutils_fsspec import FsLike, resolve_fs

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, SkipExhausted, pause, resume, skip
from zkl_aiutils_datasets.scripting import make_dataset_metadata, write_dataset_metadata
from zkl_aiutils_datasets.wrapping import SequenceDataset

_coo_file_path = "coo.bin"
_data_file_path = "data.bin.zst"


class BytesDataset(Dataset[bytes]):
    def __init__(self, fs: FsLike):
        self._fs = resolve_fs(fs)

    @cached_property
    def _samples_coo(self) -> np.ndarray:
        with self._fs.open(_coo_file_path, "rb") as fp:
            return np.frombuffer(fp.read(), dtype=np.int64)

    @cached_property
    def _data_fp(self) -> pyzstd.ZstdFile:
        return pyzstd.open(self._fs.open(_data_file_path, "rb"))

    @property
    def samples_bytes_n(self) -> SequenceDataset[int]:
        return SequenceDataset(BytesNumSequence(self))

    @property
    def total_bytes_n(self) -> int:
        return int(self._samples_coo[-1])

    def __len__(self) -> int:
        return len(self._samples_coo) - 1

    def __getitem__(self, index: int, /) -> bytes:
        head = int(self._samples_coo[index])
        tail = int(self._samples_coo[index + 1])
        self._data_fp.seek(head)
        return self._data_fp.read(tail - head)

    def __iter__(self) -> 'BytesDatasetIterator':
        return BytesDatasetIterator(self)

    def __resume__(self, state: int | None = None, /) -> 'BytesDatasetIterator':
        return BytesDatasetIterator(self, state)


class BytesNumSequence(Sequence[int]):
    def __init__(self, dataset: BytesDataset):
        self._samples_coo = dataset._samples_coo

    def __len__(self) -> int:
        return len(self._samples_coo) - 1

    def __getitem__(self, index: int) -> int:
        head = int(self._samples_coo[index])
        tail = int(self._samples_coo[index + 1])
        return int(tail - head)


class BytesDatasetIterator(DatasetIterator[bytes]):
    def __init__(self, dataset: BytesDataset, state: int | None = None):
        self._data_fp = pyzstd.open(dataset._fs.open(_data_file_path, "rb"))
        self._coo_iterator = resume(SequenceDataset(dataset._samples_coo), state)
        self._head = int(next(self._coo_iterator))

    def __next__(self) -> bytes:
        self._data_fp.seek(self._head)
        tail = int(next(self._coo_iterator))
        sample = self._data_fp.read(tail - self._head)
        self._head = tail
        return sample

    def __skip__(self, samples_n: int):
        skip(self._coo_iterator, samples_n - 1)
        try:
            self._head = int(next(self._coo_iterator))
        except StopIteration:
            raise SkipExhausted(samples_n - 1)

    def __pause__(self) -> int:
        return pause(self._coo_iterator) - 1


class BytesDatasetWriter(DatasetWriter[bytes]):
    def __init__(self,
        fs: FsLike, *,
        size_limit: int | None = None,
        count_limit: int | None = None,
        loader: Callable | type | None = BytesDataset,
    ):
        self._size_limit = size_limit
        self._count_limit = count_limit

        self._fs = resolve_fs(fs)
        self._fs.makedirs("", exist_ok=True)

        if loader is not None:
            write_dataset_metadata(fs, make_dataset_metadata(loader))

        self._sample_i = 0
        self._sample_head = 0
        self._coo_fp = self._fs.open(_coo_file_path, "wb")
        self._coo_fp.write(np.array([0], dtype=np.int64).tobytes())
        self._data_fp = pyzstd.open(self._fs.open(_data_file_path, "wb"), 'wb')

    def write(self, sample: bytes):
        if self._count_limit is not None and self._sample_i + 1 >= self._count_limit:
            raise StopIteration
        if self._size_limit is not None and self._sample_head + len(sample) >= self._size_limit:
            raise StopIteration

        tail = self._sample_head + len(sample)
        self._data_fp.write(sample)
        self._coo_fp.write(np.array([tail], dtype=np.int64).tobytes())
        self._sample_i += 1
        self._sample_head = tail

    def flush(self):
        self._coo_fp.flush()
        self._data_fp.flush()

    def close(self):
        self._coo_fp.close()
        self._data_fp.close()
